#ifndef __register_H__
#define __register_H__

#include "rendPlatform.hpp"
#include <vector>
#include <map>
#include <assert.h>
#include <stdlib.h>

namespace rend
{

namespace naked
{

// class Register - a searchable container of stack construction, implemented as a contained std::map
// WARNING:
// class relies on the fact that std::map is implemented through std::_Rb_tree, which uses std::_Rb_tree_node as node

template < typename KEY_T, typename VALUE_T >
class Register
{

#define ALLOC_SIZE(type) ((sizeof(type) + 3) & ~3)

	struct Storage
	{
		void*	 bits;
		unsigned capacity;
		unsigned size;
	};
	
	static Storage *s_storage;

	template < typename T >
	class Allocator
	{
	public:
		typedef size_t    size_type;
		typedef ptrdiff_t difference_type;
		typedef T*        pointer;
		typedef const T*  const_pointer;
		typedef T&        reference;
		typedef const T&  const_reference;
		typedef T         value_type;

		template < typename U >
		struct rebind
		{ typedef Allocator<U> other; };

		Allocator() {}

		Allocator(const Allocator&) {}

		template < typename U >
		Allocator(const Allocator<U>&) {}

		~Allocator() {}

		pointer
		address(reference x) const { return &x; }

		const_pointer
		address(const_reference x) const { return &x; }

		pointer
		allocate(size_type n, const void * = 0)
		{
			assert(s_storage);
			const size_type new_size = s_storage->size + n;

			if (!s_storage->bits || new_size > s_storage->capacity)
				std::__throw_bad_alloc();

			T* p = reinterpret_cast< T* >((int8*) s_storage->bits + s_storage->size * ALLOC_SIZE(T));
			s_storage->size = new_size;

			return p;
		}

		void
		deallocate(void* p, size_type) {}

		void
		construct(pointer p, const T& val)
		{
			::new(p) T(val);
		}

		void
		destroy(pointer p)
		{
			p->~T();
		}

		size_type
		max_size() const
		{
			assert(s_storage);
			return s_storage->capacity;
		}
	};

	typedef std::pair< const KEY_T, VALUE_T >	Pair;
	typedef std::_Rb_tree_node< Pair >			IndexNode;

	typedef std::map< KEY_T, VALUE_T,
					  std::less< KEY_T >,
					  Allocator< Pair > >		Index;

public:

	typedef KEY_T								Key;
	typedef VALUE_T								Value;

	// construct with specified element capacity
	Register(const unsigned capacity);

	~Register();

	// register's current size
	unsigned
	size() const;

	// register's capacity
	unsigned
	capacity() const;

	// discard last 'current_size - new_size' elements; if new_size is greater-or-equal to current_size
	// function is a nop; invokes optional visitor on the values of those elements being discarded
	void
	discard(const unsigned new_size, void (*discard_func)(Value&) = 0);

	// insert element with specified key; introduces a new element, or updates an already existing one; returns null at failure
	Value*
	insert(const Key& key);

	// seek element with specified key; returns null at failure
	const Value*
	seek(const Key& key) const;

private:

	Register() {}
	Register(const Register< KEY_T, VALUE_T >&) {}

	Register< KEY_T, VALUE_T >&
	operator =(const Register< KEY_T, VALUE_T >&) {}

	Storage		m_storage;
	Index		m_index;
};


template < typename KEY_T, typename VALUE_T >
Register< KEY_T, VALUE_T >::Register(const unsigned capacity)
{
	m_storage.bits = malloc(capacity * ALLOC_SIZE(IndexNode));
	assert(m_storage.bits);

	m_storage.capacity = capacity;
	m_storage.size = 0;

#if defined(DEBUG)

	fprintf(stdout, "constructed Register, storage at %p, capacity of %u, %u-byte units, %u-byte payload\n",
		&m_storage.bits, m_storage.capacity, (unsigned) ALLOC_SIZE(IndexNode), (unsigned) sizeof(VALUE_T));

#endif
}


template < typename KEY_T, typename VALUE_T >
Register< KEY_T, VALUE_T >::~Register()
{
	discard(0); // as we cannot pass a discard visitor here, values of POD types must have already been propely discarded

	free(m_storage.bits);

#if defined(DEBUG)

	fprintf(stdout, "destroyed Register, storage at %p\n", &m_storage.bits);

#endif
}


template < typename KEY_T, typename VALUE_T >
unsigned
Register< KEY_T, VALUE_T >::size() const
{
	return m_storage.size;
}


template < typename KEY_T, typename VALUE_T >
unsigned
Register< KEY_T, VALUE_T >::capacity() const
{
	return m_storage.capacity;
}


template < typename KEY_T, typename VALUE_T >
void
Register< KEY_T, VALUE_T >::discard(const unsigned new_size, void (*discard_func)(Value&))
{
	if (new_size < m_storage.size)
	{
		for (unsigned i = new_size; i < m_storage.size; ++i)
		{
			IndexNode* node = reinterpret_cast< IndexNode* >((int8*) m_storage.bits + i * ALLOC_SIZE(IndexNode));

			if (discard_func)
				discard_func(node->_M_value_field.second);

			m_index.erase(node->_M_value_field.first);
		}

		m_storage.size = new_size;
	}
}


template < typename KEY_T, typename VALUE_T >
typename Register< KEY_T, VALUE_T >::Value*
Register< KEY_T, VALUE_T >::insert(const Key& key)
{
	typename Index::iterator it = m_index.find(key);

	if (it != m_index.end())
		return &it->second;

	assert(m_storage.size < m_storage.capacity);

	if (m_storage.size == m_storage.capacity)
		return 0;

	s_storage = &m_storage;
	std::pair< typename Index::iterator, bool > r = m_index.insert(typename Index::value_type(key, Value()));
	s_storage = 0;

	assert(r.second);

	return &r.first->second;
}


template < typename KEY_T, typename VALUE_T >
const typename Register< KEY_T, VALUE_T >::Value*
Register< KEY_T, VALUE_T >::seek(const Key& key) const
{
	typename Index::const_iterator it = m_index.find(key);

	if (it != m_index.end())
		return &it->second;

	return 0;
}

} // namespace naked
} // namespace rend

#endif // __register_H__
